"""
使用机器学习模型识别主力建仓模式
特征工程包括：
1. 技术指标（MACD, KDJ, RSI等）
2. 量价关系指标
3. 趋势指标
4. 波动率指标
5. 资金流向指标
"""
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from sqlalchemy import select, and_
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import IsolationForest, RandomForestClassifier
from sklearn.model_selection import train_test_split
import talib
from models.stock_models import StockDailyData, StockRealtimeData
from database.db_engine import SessionLocal

def calculate_technical_indicators(df):
    """计算技术指标"""
    # 确保数据按时间排序
    df = df.sort_values('trade_date')
    
    # 1. 趋势指标
    # MACD
    df['macd'], df['macd_signal'], df['macd_hist'] = talib.MACD(
        df['close'], fastperiod=12, slowperiod=26, signalperiod=9
    )
    
    # 2. 动量指标
    # RSI
    df['rsi'] = talib.RSI(df['close'], timeperiod=14)
    
    # KDJ
    df['slowk'], df['slowd'] = talib.STOCH(
        df['high'], df['low'], df['close'],
        fastk_period=9, slowk_period=3, slowk_matype=0,
        slowd_period=3, slowd_matype=0
    )
    
    # 3. 波动率指标
    # Bollinger Bands
    df['bb_upper'], df['bb_middle'], df['bb_lower'] = talib.BBANDS(
        df['close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0
    )
    
    # ATR - Average True Range
    df['atr'] = talib.ATR(df['high'], df['low'], df['close'], timeperiod=14)
    
    # 4. 量价关系指标
    df['obv'] = talib.OBV(df['close'], df['volume'])
    df['ad'] = talib.AD(df['high'], df['low'], df['close'], df['volume'])
    
    # 5. 自定义指标
    # 价格动量
    df['price_momentum'] = df['close'].pct_change(5)
    # 成交量动量
    df['volume_momentum'] = df['volume'].pct_change(5)
    # 量价背离指标
    df['price_volume_divergence'] = df['price_momentum'] - df['volume_momentum']
    
    # 6. 移动平均线
    df['ma5'] = talib.MA(df['close'], timeperiod=5)
    df['ma10'] = talib.MA(df['close'], timeperiod=10)
    df['ma20'] = talib.MA(df['close'], timeperiod=20)
    df['ma60'] = talib.MA(df['close'], timeperiod=60)
    
    return df

def prepare_features(df):
    """准备特征数据"""
    feature_columns = [
        'macd', 'macd_signal', 'macd_hist',
        'rsi', 'slowk', 'slowd',
        'bb_upper', 'bb_middle', 'bb_lower', 'atr',
        'obv', 'ad',
        'price_momentum', 'volume_momentum', 'price_volume_divergence',
        'turnover_rate'
    ]
    
    # 计算技术指标
    df = calculate_technical_indicators(df)
    
    # 处理缺失值
    df = df.fillna(method='ffill')
    
    # 标准化特征
    scaler = StandardScaler()
    features = scaler.fit_transform(df[feature_columns])
    
    return pd.DataFrame(features, columns=feature_columns, index=df.index)

def train_anomaly_detector(features):
    """训练异常检测模型"""
    # 使用IsolationForest进行异常检测
    model = IsolationForest(
        n_estimators=100,
        contamination=0.1,  # 假设10%的样本是异常值（主力建仓）
        random_state=42
    )
    model.fit(features)
    return model

def detect_main_force_patterns(days=60):
    """
    使用机器学习模型检测主力建仓模式
    
    Args:
        days (int): 分析的历史数据天数
    
    Returns:
        list: 检测到的主力建仓股票列表
    """
    end_date = datetime.now().date()
    start_date = end_date - timedelta(days=days)
    
    session = SessionLocal()
    try:
        # 获取历史数据
        query = select(StockDailyData).where(
            and_(
                StockDailyData.trade_date >= start_date,
                StockDailyData.trade_date <= end_date
            )
        )
        results = session.execute(query).scalars().all()
        
        # 转换为DataFrame
        df = pd.DataFrame([{
            'ts_code': r.ts_code,
            'trade_date': r.trade_date,
            'open': r.open,
            'high': r.high,
            'low': r.low,
            'close': r.close,
            'volume': r.volume,
            'amount': r.amount,
            'turnover_rate': r.turnover_rate
        } for r in results])
        
        if df.empty:
            return []
        
        result_stocks = []
        # 按股票代码分组分析
        for ts_code, group in df.groupby('ts_code'):
            if len(group) < 60:  # 确保有足够的数据
                continue
                
            # 准备特征
            features = prepare_features(group)
            
            # 训练模型并预测
            model = train_anomaly_detector(features)
            predictions = model.predict(features)
            
            # 获取最近的预测结果
            recent_predictions = predictions[-5:]  # 最近5天的预测
            
            # 如果最近5天中有3天以上被判定为异常（-1），则认为可能存在主力建仓
            if sum(recent_predictions == -1) >= 3:
                # 计算置信度分数
                confidence_score = model.score_samples(features.iloc[-5:])
                avg_confidence = np.mean(np.abs(confidence_score))
                
                result_stocks.append({
                    'ts_code': ts_code,
                    'last_date': group['trade_date'].iloc[-1],
                    'confidence': avg_confidence,
                    'abnormal_days': sum(recent_predictions == -1),
                    'features': {
                        'rsi': features['rsi'].iloc[-1],
                        'macd': features['macd'].iloc[-1],
                        'volume_momentum': features['volume_momentum'].iloc[-1]
                    }
                })
        
        # 按置信度排序
        result_stocks.sort(key=lambda x: x['confidence'], reverse=True)
        return result_stocks
    
    finally:
        session.close()

def print_ml_analysis_results():
    """打印机器学习分析结果"""
    results = detect_main_force_patterns()
    if not results:
        print("没有发现疑似主力建仓的股票")
        return
    
    print(f"\n通过机器学习模型发现 {len(results)} 只疑似主力建仓的股票：")
    print("=" * 70)
    for stock in results:
        print(f"股票代码: {stock['ts_code']}")
        print(f"最新交易日: {stock['last_date']}")
        print(f"置信度: {stock['confidence']:.4f}")
        print(f"异常天数: {stock['abnormal_days']}/5")
        print("技术指标:")
        print(f"  RSI: {stock['features']['rsi']:.2f}")
        print(f"  MACD: {stock['features']['macd']:.2f}")
        print(f"  成交量动量: {stock['features']['volume_momentum']:.2f}")
        print("-" * 70)

if __name__ == "__main__":
    print_ml_analysis_results() 